import geopandas as gpd
import matplotlib.pyplot as plt
import numpy as np
import pandas as p

# State shape file from https://www.census.gov/geographies/mapping-files/time-series/geo/carto-boundary-file.html

shape = 'cb_2018_us_state_500k 2/cb_2018_us_state_500k.shp'
# shape = r"C:\Users\Connor\Box\CMU\Data\Map Data\cb_2018_us_state_500k 2/cb_2018_us_state_500k.shp"
map = gpd.read_file(shape)

contigRows = (map['STUSPS']!='HI') & (map['STUSPS']!='AK')

map = map[contigRows]

cleanData = 'FHWACleanV4.csv'

cleanData = p.read_csv(cleanData)
cleanData = cleanData[cleanData['Year']==1995]
cleanData = cleanData.replace('Dist. of Col.', 'District of Columbia')

map = map.merge(cleanData, left_on='NAME', right_on='State')


res = 400
lw = 0.5
x = 10
gr = (1+np.sqrt(5))/2
y = x/(gr*2)

fig, ax = plt.subplots(1,2, figsize=(x,y))



map.plot(ax=ax[0], color='grey', edgecolor='k', linewidth=lw)
map.plot(ax=ax[0], column='StartedSafety', edgecolor='k', linewidth=lw, cmap='hot',
                   legend=True, legend_kwds={'orientation':'horizontal'})

plt.sca(ax[0])
plt.yticks([])
plt.xticks([])
title = "Contiguous US \n by Year Safety Inspections Began"
title = "Year Safety Inspections Began"
plt.title(title)

map.plot(ax=ax[1], color='grey', edgecolor='k', linewidth=lw)
map.plot(ax=ax[1], column='EndedSafety', linewidth=lw, edgecolor='k', cmap='bwr',
                   legend=True, legend_kwds={'orientation':'horizontal'})

plt.sca(ax[1])
plt.yticks([])
plt.xticks([])
title = "Contiguous US \n by Year Safety Inspections Ended"
title = "Year Safety Inspections Ended"
plt.title(title)

file = 'Plots/Cont48InspectionBounds.png'
file = 'Plots/Cont48InspectionBounds.pdf'
plt.savefig(file, bbox_inches="tight", dpi=res)

res = 1000
x = 4
gr = (1+np.sqrt(5))/2
y = x/(gr)

fig, ax = plt.subplots(1,1, figsize=(x,y))
map.plot(ax=ax, color='grey', edgecolor='k', linewidth=lw)
map.plot(ax=ax, column='EndedSafety', linewidth=lw, edgecolor='k', cmap='bwr',
                   legend=True, legend_kwds={'orientation':'horizontal', 'pad':0.05})

plt.sca(ax)
# plt.colorbar(ax, size=15)
plt.yticks([])
plt.xticks([])
title = "Year Safety Inspections Ended"
plt.title(title, fontsize=12)
file = 'Plots/Cont48EndInspectionBounds.png'
file = 'Plots/Cont48EndInspectionBounds.pdf'
# plt.show()
plt.savefig(file, bbox_inches="tight", dpi=res)